#include "stdafx.h"
#include "../aclib/HexParser.h"
#include "../aclib/BlockData.h"

#include "../aclib/OFB.h"

#include "MCryptor.h"
#include "MLastBlockHandler.h"
#include "MBlockData.h"

namespace AC = AppliedCryptography;

using ::testing::Return;
using ::testing::ReturnRef;
using ::testing::_;
using ::testing::Eq;
using ::testing::Invoke;
using ::testing::Ref;
using ::testing::AtLeast;

template<unsigned int TBlockSize>
void MCryptorProcess2(const AC::Block<TBlockSize>& plaintext, AC::Block<TBlockSize>& chyper)
{
	AC::HexParser Ox;

	AC::Block<3> iv(Ox("102030"));
	AC::Block<3> bc1(Ox("202020")), bc2(Ox("404040"));
	AC::Block<3> error(Ox("666999"));

	if(plaintext == iv)
	{
		chyper = bc1;
		return;
	};

	if(plaintext == bc1)
	{
		chyper = bc2;
		return;
	};

	chyper = error;
};

TEST(OutputFeedback, Encryption)
{
	AC::HexParser Ox;

	MLastBlockEncryption<4, 3> steal;
	MEncryptor<4, 3> cryptor;
	AC::OFB<4, 3, AC::Encryptor<4, 3>, AC::LastBlockEncryption<4, 3>>::Encryption c;

	AC::Block<3> iv(Ox("102030"));
	MBlockData<3> input;
	MBlockData<3> output;

	// Blocks returned from input BlockData object
	AC::Block<3> bp1(Ox("111213")), bp2(Ox("212223")), bp3(Ox("313200"));
	EXPECT_CALL(input, Blocks())
		.WillRepeatedly(Return(3));
	EXPECT_CALL(input, doConstOperatorAccess(0))
		.Times(AtLeast(1))
		.WillRepeatedly(Return(bp1));
	EXPECT_CALL(input, doConstOperatorAccess(1))
		.Times(AtLeast(1))
		.WillRepeatedly(Return(bp2));
	EXPECT_CALL(input, doConstOperatorAccess(2))
		.Times(AtLeast(1))
		.WillRepeatedly(Return(bp3));
	EXPECT_CALL(input, LastBlockLength())
		.WillRepeatedly(Return(2));

	// Block passed to encryption
	AC::Block<3> bc1(Ox("202020")), bc2(Ox("404040"));
	EXPECT_CALL(cryptor, Process(Eq(iv), _))
		.Times(1)
		.WillOnce(Invoke(MCryptorProcess2<3>));
	EXPECT_CALL(cryptor, Process(Eq(bc1), _))
		.Times(1)
		.WillOnce(Invoke(MCryptorProcess2<3>));
	/*
	// Skiped as this part is handled by mock LBH
	EXPECT_CALL(cryptor, Process(Eq(bc2), _))
		.Times(1)
		.WillOnce(Invoke(MCryptorProcess<3>));
	*/

	// Encrypted blocks passed to output BlockData object
	AC::Block<3> bcc3(Ox("b1b2b3"));
	AC::Block<3> bx1(Ox("313233")), bx2(Ox("616263"));
	EXPECT_CALL(output, Add(Eq(bx1), 3))
		.Times(1);
	EXPECT_CALL(output, Add(Eq(bx2), 3))
		.Times(1);
	EXPECT_CALL(output, Add(Eq(bcc3), 2))
		.Times(1);

	EXPECT_CALL(steal, Process(_, Eq(bp3), 2, Eq(bc2)))
		.Times(1)
		.WillRepeatedly(Return(bcc3));

	c.Process(cryptor, steal, iv, input, output);
};

TEST(OutputFeedback, EncryptionOneBlock)
{
	AC::HexParser Ox;

	MLastBlockEncryption<4, 3> steal;
	MEncryptor<4, 3> cryptor;
	AC::OFB<4, 3, AC::Encryptor<4, 3>, AC::LastBlockEncryption<4, 3>>::Encryption c;

	AC::Block<3> iv(Ox("102030"));
	MBlockData<3> input;
	MBlockData<3> output;

	// Blocks returned from input BlockData object
	AC::Block<3> bp1(Ox("111213"));
	EXPECT_CALL(input, Blocks())
		.WillRepeatedly(Return(1));
	EXPECT_CALL(input, doConstOperatorAccess(0))
		.Times(AtLeast(1))
		.WillRepeatedly(Return(bp1));
	EXPECT_CALL(input, LastBlockLength())
		.WillRepeatedly(Return(3));

	// Block passed to encryption
	AC::Block<3> bc1(Ox("202020"));
	EXPECT_CALL(cryptor, Process(Eq(iv), _))
		.Times(1)
		.WillOnce(Invoke(MCryptorProcess2<3>));

	// Encrypted blocks passed to output BlockData object
	AC::Block<3> bx1(Ox("313233"));
	EXPECT_CALL(output, Add(Eq(bx1), 3))
		.Times(1);

	EXPECT_CALL(steal, Process(_, Eq(bp1), 3, Eq(bc1)))
		.Times(1)
		.WillRepeatedly(Return(bx1));

	c.Process(cryptor, steal, iv, input, output);
};

TEST(OutputFeedback, Decryption)
{
	AC::HexParser Ox;

	MLastBlockDecryption<4, 3> steal;
	MDecryptor<4, 3> cryptor;
	AC::OFB<4, 3, AC::Decryptor<4, 3>, AC::LastBlockDecryption<4, 3>>::Decryption c;

	AC::Block<3> iv(Ox("102030"));
	MBlockData<3> input;
	MBlockData<3> output;

	// Blocks returned from input BlockData object
	AC::Block<3> bcc3(Ox("b1b2b3"));
	AC::Block<3> bx1(Ox("313233")), bx2(Ox("616263"));
	EXPECT_CALL(input, Blocks())
		.WillRepeatedly(Return(3));
	EXPECT_CALL(input, doConstOperatorAccess(0))
		.Times(AtLeast(1))
		.WillRepeatedly(Return(bx1));
	EXPECT_CALL(input, doConstOperatorAccess(1))
		.Times(AtLeast(1))
		.WillRepeatedly(Return(bx2));
	EXPECT_CALL(input, doConstOperatorAccess(2))
		.Times(AtLeast(1))
		.WillRepeatedly(Return(bcc3));
	EXPECT_CALL(input, LastBlockLength())
		.WillRepeatedly(Return(2));

	// Block passed to encryption
	AC::Block<3> bc1(Ox("202020")), bc2(Ox("404040"));
	EXPECT_CALL(cryptor, Process(Eq(iv), _))
		.Times(1)
		.WillOnce(Invoke(MCryptorProcess2<3>));
	EXPECT_CALL(cryptor, Process(Eq(bc1), _))
		.Times(1)
		.WillOnce(Invoke(MCryptorProcess2<3>));
	/*
	// Skiped as this part is handled by mock LBH
	EXPECT_CALL(cryptor, Process(Eq(bc2), _))
		.Times(1)
		.WillOnce(Invoke(MCryptorProcess<3>));
	*/

	// Encrypted blocks passed to output BlockData object
	AC::Block<3> bp1(Ox("111213")), bp2(Ox("212223")), bp3(Ox("313200"));
	EXPECT_CALL(output, Add(Eq(bp1), 3))
		.Times(1);
	EXPECT_CALL(output, Add(Eq(bp2), 3))
		.Times(1);
	EXPECT_CALL(output, Add(Eq(bp3), 2))
		.Times(1);

	EXPECT_CALL(steal, Process(_, Eq(bcc3), 2, Eq(bc2)))
		.Times(1)
		.WillRepeatedly(Return(bp3));

	c.Process(cryptor, steal, iv, input, output);
};

TEST(OutputFeedback, DecryptionOneBlock)
{
	AC::HexParser Ox;

	MLastBlockDecryption<4, 3> steal;
	MDecryptor<4, 3> cryptor;
	AC::OFB<4, 3, AC::Decryptor<4, 3>, AC::LastBlockDecryption<4, 3>>::Decryption c;

	AC::Block<3> iv(Ox("102030"));
	MBlockData<3> input;
	MBlockData<3> output;

	// Blocks returned from input BlockData object
	AC::Block<3> bx1(Ox("313233"));
	EXPECT_CALL(input, Blocks())
		.WillRepeatedly(Return(1));
	EXPECT_CALL(input, doConstOperatorAccess(0))
		.Times(AtLeast(1))
		.WillRepeatedly(Return(bx1));
	EXPECT_CALL(input, LastBlockLength())
		.WillRepeatedly(Return(3));

	// Block passed to encryption
	AC::Block<3> bc1(Ox("202020"));
	EXPECT_CALL(cryptor, Process(Eq(iv), _))
		.Times(1)
		.WillOnce(Invoke(MCryptorProcess2<3>));

	// Encrypted blocks passed to output BlockData object
	AC::Block<3> bp1(Ox("111213"));
	EXPECT_CALL(output, Add(Eq(bp1), 3))
		.Times(1);

	EXPECT_CALL(steal, Process(_, Eq(bx1), 3, Eq(bc1)))
		.Times(1)
		.WillRepeatedly(Return(bp1));

	c.Process(cryptor, steal, iv, input, output);
};
